// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

#pragma once

#include <nebula/util/concurrent_map.h>

#include <aries/parquet/encryption/encryption.h>
#include <aries/parquet/encryption/file_system_key_material_store.h>
#include <aries/parquet/encryption/key_material.h>
#include <aries/parquet/encryption/key_toolkit.h>
#include <aries/parquet/encryption/key_toolkit_internal.h>
#include <aries/parquet/encryption/kms_client.h>
#include <aries/parquet/platform.h>

namespace parquet::encryption {
    // This class will retrieve the key from "key metadata", following these steps:
    // 1. Parse "key metadata" (see structure in KeyMetadata class).
    // 2. Retrieve "key material" which can be stored inside or outside "key metadata".
    // 3. Unwrap the "data encryption key" from "key material". There are 2 modes:
    // 3.1. single wrapping: decrypt the wrapped "data encryption key" directly with "master
    // encryption key" 3.2. double wrapping: 2 steps: 3.2.1. "key encryption key" is decrypted
    // with "master encryption key" 3.2.2. "data encryption key" is decrypted with the above
    // "key encryption key"
    class PARQUET_EXPORT FileKeyUnwrapper : public DecryptionKeyRetriever {
    public:
        /// key_toolkit and kms_connection_config is to get KmsClient from cache or create
        /// KmsClient if it's not in the cache yet. cache_entry_lifetime_seconds is life time of
        /// KmsClient in the cache.
        /// If the file uses external key material then the Parquet file path and file
        /// system must be specified.
        FileKeyUnwrapper(std::shared_ptr<KeyToolkit> key_toolkit,
                         const KmsConnectionConfig &kms_connection_config,
                         double cache_lifetime_seconds, const std::string &file_path = "",
                         const std::shared_ptr<alioth::FileSystem> &file_system = nullptr);

        /// Constructor overload that takes a raw pointer to the KeyToolkit
        FileKeyUnwrapper(KeyToolkit *key_toolkit,
                         const KmsConnectionConfig &kms_connection_config,
                         double cache_lifetime_seconds, const std::string &file_path = "",
                         const std::shared_ptr<alioth::FileSystem> &file_system = nullptr);

        /// Constructor overload that takes a raw pointer to the KeyToolkit and
        /// accepts an existing key_material_store rather than using
        /// the file path and file system to create one when needed.
        FileKeyUnwrapper(KeyToolkit *key_toolkit,
                         const KmsConnectionConfig &kms_connection_config,
                         double cache_lifetime_seconds,
                         std::shared_ptr<FileKeyMaterialStore> key_material_store);

        /// Get the data key from key metadata
        std::string GetKey(const std::string &key_metadata) override;

        /// Get the data key along with the master key id from key material
        KeyWithMasterId GetDataEncryptionKey(const KeyMaterial &key_material);

    private:
        FileKeyUnwrapper(std::shared_ptr<KeyToolkit> key_toolkit_owner, KeyToolkit *key_toolkit,
                         const KmsConnectionConfig &kms_connection_config,
                         double cache_lifetime_seconds,
                         std::shared_ptr<FileKeyMaterialStore> key_material_store,
                         const std::string &file_path,
                         const std::shared_ptr<alioth::FileSystem> &file_system);

        std::shared_ptr<KmsClient> GetKmsClientFromConfigOrKeyMaterial(
            const KeyMaterial &key_material);

        /// A map of Key Encryption Key (KEK) ID -> KEK bytes, for the current token
        std::shared_ptr<::nebula::util::ConcurrentMap<std::string, std::string> > kek_per_kek_id_;
        std::shared_ptr<KeyToolkit> key_toolkit_owner_;
        KeyToolkit *key_toolkit_;
        KmsConnectionConfig kms_connection_config_;
        const double cache_entry_lifetime_seconds_;
        std::shared_ptr<FileKeyMaterialStore> key_material_store_;
        const std::string file_path_;
        std::shared_ptr<alioth::FileSystem> file_system_;
    };
} // namespace parquet::encryption
